from metadrive.engine import get_global_config
from metadrive.engine.core.manual_controller import KeyboardController, SteeringWheelController
from metadrive.policy.manual_control_policy import TakeoverPolicy



class MyKeyboardController(KeyboardController):
    # Update Parameters
    STEERING_INCREMENT = 0.05
    STEERING_DECAY = 0.5

    THROTTLE_INCREMENT = 0.5
    THROTTLE_DECAY = 1

    BRAKE_INCREMENT = 0.5
    BRAKE_DECAY = 1

class MyTakeoverPolicy(TakeoverPolicy):
    """
    Record the takeover signal
    """
    def __init__(self, obj, seed):
        super(TakeoverPolicy, self).__init__(obj, seed)
        config = get_global_config()
        if config["manual_control"] and config["use_render"]:
            if config["controller"] == "joystick":
                self.controller = SteeringWheelController()
            elif config["controller"] == "keyboard":
                self.controller = MyKeyboardController(False)
            else:
                raise ValueError("Unknown Policy: {}".format(config["controller"]))
        self.takeover = False
